package com.zhuodewen.ai.config.ai;

import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiImageModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import lombok.Getter;
import lombok.Setter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

@Component
@Getter
@Setter
public class LangChain4JProvider {

    private final OpenAiChatModel          openAiChatModel;
    private final OpenAiStreamingChatModel openAiStreamingChatModel;

    private final OllamaChatModel          ollamaChatModel;
    private final OllamaStreamingChatModel ollamaStreamingChatModel;

    private final OllamaEmbeddingModel ollamaEmbeddingModel;

    private final OpenAiImageModel     openAiImageModel;

    private LangChain4JOpenAIConfig langChain4JOpenAIConfig;
    private LangChain4JOllamaConfig langChain4JOllamaConfig;

    @Autowired
    public LangChain4JProvider(LangChain4JOpenAIConfig langChain4JOpenAIConfig, LangChain4JOllamaConfig langChain4JOllamaConfig) {
        this.langChain4JOpenAIConfig = langChain4JOpenAIConfig;
        this.langChain4JOllamaConfig = langChain4JOllamaConfig;

        this.openAiChatModel = new OpenAiChatModel.OpenAiChatModelBuilder()
                .baseUrl(langChain4JOpenAIConfig.getChatModel().getBaseUrl())
                .apiKey(langChain4JOpenAIConfig.getChatModel().getApiKey())
                .modelName(langChain4JOpenAIConfig.getChatModel().getModelName()).build();

        this.openAiStreamingChatModel = new OpenAiStreamingChatModel.OpenAiStreamingChatModelBuilder()
                .baseUrl(langChain4JOpenAIConfig.getStreamingChatModel().getBaseUrl())
                .apiKey(langChain4JOpenAIConfig.getStreamingChatModel().getApiKey())
                .modelName(langChain4JOpenAIConfig.getStreamingChatModel().getModelName())
                .build();

        this.ollamaChatModel = OllamaChatModel.builder()
                .baseUrl(langChain4JOllamaConfig.getChatModel().getBaseUrl())
                .modelName(langChain4JOllamaConfig.getChatModel().getModelName())
                .build();

        this.ollamaStreamingChatModel = OllamaStreamingChatModel.builder()
                .baseUrl(langChain4JOllamaConfig.getStreamingChatModel().getBaseUrl())
                .modelName(langChain4JOllamaConfig.getStreamingChatModel().getModelName())
                .build();

        this.ollamaEmbeddingModel = OllamaEmbeddingModel.builder()
                .baseUrl(langChain4JOllamaConfig.getEmbeddingModel().getBaseUrl())
                .modelName(langChain4JOllamaConfig.getEmbeddingModel().getModelName())
                .build();

        this.openAiImageModel = OpenAiImageModel.builder()
                .baseUrl(langChain4JOpenAIConfig.getImageModel().getBaseUrl())
                .apiKey(langChain4JOpenAIConfig.getImageModel().getApiKey())
                .modelName(langChain4JOpenAIConfig.getImageModel().getModelName())
                .build();

    }

}
